package com.utils.rsa;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.HashMap;

public class RSAUtils {
    private static final String PUBLIC_KEY = "RSAPublicKey";
    private static final String PRIVATE_KEY = "RSAPrivateKey";
    private static final String ALGORITHM_RSA = "RSA";
    private static final String ALGORITHM_MD5_RSA = "MD5withRSA";
    private static final int KEY_SIZE = 1024;
    private static final int MAX_DECRYPT_BLOCK = 128;// MAX_DECRYPT_BLOCK = KEY_SIZE / 8;
    private static final int MAX_ENCRYPT_BLOCK = 117;

    public static void main(String[] args) throws GeneralSecurityException, IOException {
        HashMap<String, Key> hashMap = genKeyPair();
        String publicKey = getPublicKeyBStr(hashMap);
        String privateKey = getPrivateKeyBStr(hashMap);
        System.out.println(publicKey);
        System.out.println(privateKey);

        byte[] data = "签名".getBytes(StandardCharsets.UTF_8);
        String sign = sign(privateKey, data);
        System.out.println(sign);

        boolean verify = verify(publicKey, data, sign);
        System.out.println(verify);
        System.out.println("___________------___________");
        byte[] data1 = "需要的加密内容".getBytes(StandardCharsets.UTF_8);
        System.out.println(Arrays.toString(data1));
        byte[] encryptedData = encryptByPublicKey(publicKey, data1);
        System.out.println(Arrays.toString(encryptedData));
        byte[] bytes = decryptByPrivateKey(privateKey, encryptedData);
        System.out.println(Arrays.toString(bytes));
        System.out.println(new String(bytes));
        System.out.println("___________------___________");

        System.out.println(hashMap);
    }

    public static byte[] decryptByPrivateKey(String privateKey, byte[] encryptedData) throws GeneralSecurityException, IOException {
        PrivateKey privateK = getPrivateKey(privateKey);
        return decryptByPrivateKey(privateK, encryptedData);
    }

    public static byte[] decryptByPrivateKey(PrivateKey privateKey, byte[] encryptedData) throws GeneralSecurityException, IOException {
        return doFinal(encryptedData, Cipher.DECRYPT_MODE, MAX_DECRYPT_BLOCK, privateKey);
    }

    public static byte[] encryptByPublicKey(String publicKey, byte[] data) throws GeneralSecurityException, IOException {
        PublicKey publicK = getPublicKey(publicKey);
        return encryptByPublicKey(publicK, data);
    }

    public static byte[] encryptByPublicKey(PublicKey publicKey, byte[] data) throws GeneralSecurityException, IOException {
        return doFinal(data, Cipher.ENCRYPT_MODE, MAX_ENCRYPT_BLOCK, publicKey);
    }

    private static byte[] doFinal(byte[] encryptedData, int cipherMode, int maxBlock, Key key) throws GeneralSecurityException, IOException {
        Cipher cipher = Cipher.getInstance(ALGORITHM_RSA);
        cipher.init(cipherMode, key);
        int length = encryptedData.length;
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        for (int i = 0, offset = 0; length - offset > 0; ++i, offset = i * maxBlock) {
            byte[] cache;
            if (length - offset > maxBlock) {
                cache = cipher.doFinal(encryptedData, offset, maxBlock);
            } else {
                cache = cipher.doFinal(encryptedData, offset, length - offset);
            }
            outputStream.write(cache, 0, cache.length);
        }
        byte[] bytes = outputStream.toByteArray();
        outputStream.close();
        return bytes;
    }

    public static boolean verify(String publicKey, byte[] data, String sign) throws GeneralSecurityException {
        PublicKey generatePublic = getPublicKey(publicKey);
        Signature signature = Signature.getInstance(ALGORITHM_MD5_RSA);
        signature.initVerify(generatePublic);
        signature.update(data);
        return signature.verify(Base64.decodeBase64(sign));
    }

    public static String sign(String privateKey, byte[] data) throws GeneralSecurityException {
        PrivateKey generatePrivate = getPrivateKey(privateKey);
        Signature signature = Signature.getInstance(ALGORITHM_MD5_RSA);
        signature.initSign(generatePrivate);
        signature.update(data);
        return Base64.encodeBase64String(signature.sign());
    }

    public static PublicKey getPublicKey(String publicKey) throws GeneralSecurityException {
        byte[] bytes = Base64.decodeBase64(publicKey);
        X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(bytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
        return keyFactory.generatePublic(x509EncodedKeySpec);
    }

    public static PrivateKey getPrivateKey(String privateKey) throws GeneralSecurityException {
        byte[] bytes = Base64.decodeBase64(privateKey);
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(bytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_RSA);
        return keyFactory.generatePrivate(pkcs8EncodedKeySpec);
    }

    public static String getPublicKeyBStr(HashMap<String, Key> hashMap) {
        byte[] encoded = hashMap.get(PUBLIC_KEY).getEncoded();
        return Base64.encodeBase64String(encoded);
    }

    public static String getPrivateKeyBStr(HashMap<String, Key> hashMap) {
        byte[] encoded = hashMap.get(PRIVATE_KEY).getEncoded();
        return Base64.encodeBase64String(encoded);
    }

    public static PublicKey getPublicKey(HashMap<String, Key> hashMap) {
        return (PublicKey) hashMap.get(PUBLIC_KEY);
    }

    public static PrivateKey getPrivateKey(HashMap<String, Key> hashMap) {
        return (PrivateKey) hashMap.get(PRIVATE_KEY);
    }

    public static HashMap<String, Key> genKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator rsa = KeyPairGenerator.getInstance(ALGORITHM_RSA);
        rsa.initialize(KEY_SIZE);
        KeyPair keyPair = rsa.generateKeyPair();
        RSAPublicKey rsaPublicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) keyPair.getPrivate();
        HashMap<String, Key> hashMap = new HashMap<>(2);
        hashMap.put(PUBLIC_KEY, rsaPublicKey);
        hashMap.put(PRIVATE_KEY, rsaPrivateKey);
        return hashMap;
    }
}
